"""
This provides an example of simulating coalescent trees for a simple
HIV model. See equations 5.1-5.6 of Frost & Volz, PTRSB 2013.

Modify __main__ clause to change parameters. 
"""

from pylab import *
from scipy.integrate import odeint
import colgem2
import pdb


class HIVModel:
	""" 
	This class solves ODE and outputs timeseries for coalescent model
	"""
	
	#~ default parameters
	beta1 = .01
	beta2 = .001
	alpha = 1.
	gamma = 1/9.
	mu = 1/40.
	c = 100.
	N = 10**4.
	
	def _F(self, x):
		#~ birth rate (ie transmissions)
		s, i1, i2 = x
		return array( [[self.beta1 * self.c * i1 * s / self.N, 0.],
		[self.beta2 * self.c * i2 * s / self.N, 0.]] )
	#
	
	def _G(self, x):
		#~ migration rate (ie stage progression)
		s, i1, i2 = x
		return array([[0, self.alpha * i1],
		[0., 0.]])
	#
	
	def xdot(self, x, t):
		#~ calculates time derivative of state variables
		s, i1, i2 = x
		F = self._F(x)
		G = self._G(x)
		ds = -sum(F) - self.mu * s
		di1 = sum(F) - sum(G) - self.mu * i1
		di2 = sum(G) - self.mu * i2 - self.gamma * i2
		ds -= (ds + di1 + di2) # keeps N constant
		return (ds, di1, di2)
	#
	
	def __init__(self, t=None, x0 = None, **kwargs):
		'''
		t array (float) : times at which to solve ODEs
		x0 array(float) : initial conditions (S, I1, I2)
		parameters may be changed by passing dictionary kwargs
		'''
		self.__dict__.update(**kwargs)
		if not x0:
			x0 = array([self.N-1, 1., 0.]) #initially one acute infection
		if not t:
			t = linspace(0., 40., 100) #0..40 years
		#
		self.t = t
		x = self.x = odeint(self.xdot, x0, t)
		self.s, self.i1, self.i2 = x.T
	#
	def FGY(self):
		#generates births, migrations, and population size timeseries
		#~ uses approximation that number of events in small interval h
		#~ is h * rate
		h = self.t[1]
		Y = c_[self.i1, self.i2]
		F = [ h * self._F(xx) for xx in self.x[1:]]
		G = [h * self._G(xx) for xx in self.x[1:]]
		return F, G, Y
	#
#

if __name__=='__main__':
	#~ make some trees
	ntrees=10
	n = 100 #sample size
	
	
	m = HIVModel() #solve the model
	F, G, Y = m.FGY() #births, migration, and population size
	fgy = colgem2.FGY(m.t, F, G, Y)
	
	taxa = ['_%i_' % i for i in range(n)] # the name of each taxon
	sampleTimes = dict(zip( taxa, [m.t[-1]]*n)) # homochronous sample at last simulation time
	sampleStates = dict(zip(taxa, [eye(2)[0]]*int(.1*n) + [eye(2)[1]]*int(.9*n) ) ) # 10pc of sample will be acute
	
	outnwks = list()
	for i in range(ntrees):
		trees, nwks, A, daf = colgem2.simulate_coalescent(sampleTimes, sampleStates, fgy, singleMRCA = True)
		#~ pdb.set_trace()
		outnwks.append(nwks[0])
	#
	
	''' 
	#uncomment to visually inspect trees
	from ete2 import *
	trees = list()
	for nwk in outnwks:
		trees.append(Tree(nwk))
		#~ may then browse trees, eg trees[0].show()
	'''
	
	open('trees.nwk', 'w').writelines(outnwks)
#
